from flask import Flask,request,jsonify
import flask
import mmap
import os,sys
import cv2
import json
import time
import copy
import numpy as np
from types import SimpleNamespace
import pprint
LOCAL = False

handian_CAMS = ['正面焊点','焊点上','焊点下']
gap_CAMS = ['间隙','风叶']
ALL_CAMS = ['卡扣',
            '插线',
            'ZM','BM-UP','BM-DOWN',
            'Gap1','Gap2',
            '弯管']

# ================ ================
idx_tp = 0
yolo_instance = None
yolo_model_handian = None
yolo_model_gap = None

# is_handian = True

if not LOCAL:
    yolo_dir = 'D:\\ALGO\\media_line2\\yolov8'

    if yolo_dir in sys.path:
        pass
    else:
        sys.path.append(yolo_dir)

    import importlib
    import importlib.util
    yolo_instance = importlib.import_module('ultralytics')

    yolo_model_handian_weights = 'D:\\ALGO\\media_line2\\models\\line2\\hanjie\\best.pt'
    yolo_model_gap_weights = 'D:\\ALGO\\media_line2\\models\\line2\\jianxi\\best.pt'
    # ================ ================
else:
    yolo_dir = 'D:\\py\\work\\line3\\media_line2\\yolov8'
    if yolo_dir in sys.path:
        pass
    else:
        sys.path.append(yolo_dir)

    import importlib
    import importlib.util
    yolo_instance = importlib.import_module('ultralytics')

    # ================ ================

    yolo_model_handian_weights = 'D:\\py\\work\\yolov8\\runs\\detect\\train8\\weights\\best.pt'
    yolo_model_gap_weights = 'D:\\py\\work\\yolov8\\runs\\detect\\train10\\weights\\best.pt'


def get_model(model):
    model = yolo_instance.YOLO(model)
    
    def infer(im):
        results = model(im,iou=0.5,conf=0.2)
        # print(results)
       
        return results[0].boxes.data.cpu().numpy()
        
    return infer
    

def get_model_obb(model):
    model = yolo_instance.YOLO(model)
    
    def infer(im):
        results = model(im,iou=0.5,conf=0.2)
        
        # print(results[0].obb)
        return results[0].obb.xyxyxyxy.cpu().numpy(),results[0].obb.conf,results[0].obb.cls
    return infer

app = Flask(__name__)

@app.route("/")
def hello_world():
    return "<p>This is yolo algo for media line2.</p>"


# 获取内存映射图片
def get_img(name, file_size):
    m = mmap.mmap(-1, file_size, tagname=name, access=mmap.ACCESS_READ)
    # [0]:默认-1  [1] 文件大小 [2]: id就是文件的标签 [3]: 读写模式
    t = time.time()
    while True:
        try:
            
            if time.time() - t > 0.1:  # add 23.12.21 缩短等待时间
                m.close()  # 关闭后继续循环
                return None
            m.seek(0)  # 设置文件的当前位置
            
            image_data = np.frombuffer(m.read(file_size), dtype=np.uint8)  # 将数据转换为NumPy数组
            if image_data.size == 0:
                continue
            image = cv2.imdecode(image_data, cv2.IMREAD_COLOR)  # 直接使用二进制数组解码
            
            if image is None:
                continue
            
            m.close()
            return image  # 这一步相当于跳出循环

        except FileNotFoundError:
            print(f'未找到文件')
            m.close()  # 关闭后继续循环
            return None
        
# 心跳
@app.route('/heartbeat', methods=['GET'])
def heartbeat():
    return 'OK', 200

@app.route("/json", methods=['POST'])
def get_frame():
    global idx_tp, yolo_model_handian, yolo_model_gap
    is_handian = True

    #### 算法返给客户端的字典 ####
    # 出参
    json_data = {"code": "00000", 
                 "success": True, 
                 "msg": "算法调用成功", 
                 "data": {
                     "result": {
                         "total_result": "", 
                         "camera_name": "", 
                         "standard_result": [], 
                         "images": []
                         }
                     }
                }
    idx_tp += 1  # 图片处理计数
    t_once_start = time.time()  # 单词处理图片开始时间记录
    now_idx = copy.deepcopy(idx_tp)  # 本张图片处理的序号
    t_s1 = time.time()
    
    # json解析
    receivedJson = json.loads(str(request.data, encoding="utf-8"),
                              object_hook=lambda d: SimpleNamespace(**d))
    print(f'================received json================')
    pprint.pprint(receivedJson)
    
    # 发送json
    sendJson = json.loads(json.dumps(json_data),
                          object_hook=lambda d: SimpleNamespace(**d))
    
    # 相机名称设置错误
    if receivedJson.camera_data.name in handian_CAMS:
        is_handian = True
    elif receivedJson.camera_data.name in gap_CAMS:
        is_handian = False
    else:
        sendJson.success = False
        sendJson.msg = f'相机名称设置错误，请在 [{ALL_CAMS}]中选择名字'
        send = json.dumps(sendJson, default=lambda x: x.__dict__).encode('utf8')
        print(f'================send json================')
        pprint.pprint(sendJson)
        return send
        
    imgs = receivedJson.camera_data.images


    for im in imgs:
        out_result_image = {}
        out_result_image['id'] = im.id  # 内存映射的标签  这一步中传进来的id有相同的
        out_result_image['name'] = im.name  # 内存映射共享图片名字
        out_result_image['size'] = im.size  # 共享内存的图像大小
        out_result_image['coordinate'] = []
        
        # 解码图像
        t_s2 = time.time()
        ### 在内存映射中捞取相应图片  ###
        image = get_img(im.name, im.size)  # 通过内存名字 和长度 取图 ----> numpy img
        if image is None:
            sendJson.success = False
            sendJson.msg = f'算法无法在共享内存中读取图片,cam {receivedJson.camera_data.name}'
            send = json.dumps(sendJson, default=lambda x: x.__dict__).encode('utf8')
            print(f'================send json================')
            pprint.pprint(sendJson)
            return send
        t_len = round(time.time() - t_s2, 3)
        if time.time() - t_once_start > 99999:
            print(f"--@@@@[相机]{receivedJson.camera_data.name} [当前处理图片序号]:{now_idx} [图片名]:{im.name}, >>>>当前检测时间过长,直接跳出")
            break
        
        detect_model = None
        print('======== 相机 ========')
        if image is not None:
            
            if receivedJson.camera_data.name in handian_CAMS:  # 进行相机区分
                
                print(f'相机请求{receivedJson.camera_data.name}')
                detect_model = yolo_model_handian
            elif receivedJson.camera_data.name in gap_CAMS:
                
                print(f'相机请求{receivedJson.camera_data.name}')
                detect_model = yolo_model_gap

            if detect_model is None:
                sendJson.success = False
                sendJson.msg = f'算法无法找到检测模型，请检查相机名称设置或联系算法检查模型文件'
                send = json.dumps(sendJson, default=lambda x: x.__dict__).encode('utf8')
                print(f'================send json================')
                pprint.pprint(sendJson)
                return send
            
            # yolo模型
            t_s3 = time.time()
            if is_handian:
                result = detect_model(image)  # 确定是哪颗相机对应的模型
                print(f'==== ==== ==== ==== detect result ==== ==== ==== ==== ====')
                pprint.pprint(result)
                # temp_out.extend(out)  # 检测结果 out---> [rect, cls, confidence]
                # t_len = round(time.time() - t_s3, 3)
                print(f'==== ==== ==== ====')
                label_hanjie = {
                    'hanjie':0,
                    }
                end_hanjie = {
                    'hanjie': 0,
                    }
                if len(result) == 0:
                    sendJson.success = True
                    sendJson.msg = f'算法无法检测到目标'
                    sendJson.data.result.total_result = 'NG' if is_handian else 'OK'
                    sendJson.data.result.camera_name = receivedJson.camera_data.name
                    send = json.dumps(sendJson, default=lambda x: x.__dict__).encode('utf8')
                    print(f'================send json================')
                    pprint.pprint(sendJson)
                    return send
                plot_type = 'RECTANGLE'
                
                for x1, y1, x2, y2,conf, label in result:
                    if receivedJson.camera_data.name in handian_CAMS:
                        end_hanjie[list(label_hanjie.keys())[list(label_hanjie.values()).index(int(label))]] = 1
                        print(f'hanjie标签')
                    elif receivedJson.camera_data.name in gap_CAMS:
                        end_gap[list(label_gap.keys())[list(label_gap.values()).index(int(label))]] = 1
                        print(f'jianxi标签')
                        

                    out_result_image['coordinate'].append({
                                            "standard_id": getattr(receivedJson.camera_data.standards[0],'standard_id',''),
                                            "label": receivedJson.camera_data.standards[0].label,
                                            "threshold": receivedJson.camera_data.standards[0].threshold,
                                            "scores": float(conf),
                                            "algo_result": "YOU",
                                            "color": "GREEN",
                                            "list_draw_data": [f"检测项: {receivedJson.camera_data.standards[0].label}{int(label)}, 结果: '1'，阈值: {receivedJson.camera_data.standards[0].threshold}, 分值: {float(conf)}"],
                                            "type": plot_type,
                                            "list_point": [
                                                [int(x1),int(y1)],
                                                [int(x2),int(y1)],
                                                [int(x1),int(y2)],
                                                [int(x2),int(y2)],
                                            ]
                                        })
                end_hanjie = sum(end_hanjie.values())
                
                algo_result = '0'
                
                algo_result = str(end_hanjie)   
                
                max_scores = max([row[4] for row in result])
                
                total_result = 'OK' if end_hanjie>0 else 'NG'
                
                sendJson.data.result.images.append(
                    out_result_image
                )
                
                sendJson.data.result.camera_name = receivedJson.camera_data.name
                sendJson.data.result.total_result = total_result
                sendJson.data.result.standard_result.append({
                    'standard_id':getattr(receivedJson.camera_data.standards[0],'standard_id',''),
                    'label':receivedJson.camera_data.standards[0].label,
                    'type':receivedJson.camera_data.standards[0].type,
                    'standard':receivedJson.camera_data.standards[0].standard,
                    'threshold':receivedJson.camera_data.standards[0].threshold,
                    'scores':float(max_scores),
                    'algo_result':algo_result,
                    'detect_result':total_result,
                    
                })
                
                send = json.dumps(sendJson, default=lambda x: x.__dict__).encode('utf8')
                print(f'================send json================')
                pprint.pprint(sendJson)
                return send

            else:
                obb_xyxy,obb_conf,obb_cls = detect_model(image)
                print(f'==== ==== ==== ==== detect result ==== ==== ==== ==== ====')
                pprint.pprint(obb_xyxy)
                pprint.pprint(obb_conf)
                pprint.pprint(obb_cls)
                # temp_out.extend(out)  # 检测结果 out---> [rect, cls, confidence]
                t_len = round(time.time() - t_s3, 3)
                print(f'==== ==== ==== ====')

                # 算法结果处理
                label_gap = {
                    'gap1':0,
                    'gap2':1,
                    'gap3':2,
                }
                
                # 计数
                end_gap = {
                    'gap1': 0,
                    'gap2': 0,
                    'gap3': 0
                }
        
                plot_type = 'POLYGON'
                for _xyxy,_conf,_cls in zip(obb_xyxy,obb_conf,obb_cls):
                    x1,y1,x2,y2,x3,y3,x4,y4 = [elem for pair in _xyxy  for elem in pair]

                    if receivedJson.camera_data.name in handian_CAMS:
                        end_hanjie[list(label_hanjie.keys())[list(label_hanjie.values()).index(int(_cls))]] = 1
                        print(f'hanjie标签')
                    elif receivedJson.camera_data.name in gap_CAMS:
                        end_gap[list(label_gap.keys())[list(label_gap.values()).index(int(_cls))]] = 1
                        print(f'jianxi标签')
                        

                    out_result_image['coordinate'].append({
                                            "standard_id": getattr(receivedJson.camera_data.standards[0],'standard_id',''),
                                            "label": receivedJson.camera_data.standards[0].label,
                                            "threshold": receivedJson.camera_data.standards[0].threshold,
                                            "scores": float(_conf),
                                            "algo_result": "YOU",
                                            "color": "GREEN",
                                            "list_draw_data": [f"检测项: {receivedJson.camera_data.standards[0].label}{int(_cls)}, 结果: '1'，阈值: {receivedJson.camera_data.standards[0].threshold}, 分值: {float(_conf)}"],
                                            "type": plot_type,
                                            "list_point": [
                                                [int(x1),int(y1)],
                                                [int(x2),int(y2)],
                                                [int(x3),int(y3)],
                                                [int(x4),int(y4)],
                                            ]
                                        })

                end_gap = sum(end_gap.values())
                algo_result = '0'
                
                algo_result = str(end_gap)    
                
                max_scores = 1#max([row[4] for row in result])
                
                total_result = 'OK'
                sendJson.data.result.images.append(
                    out_result_image
                )
            
                sendJson.data.result.camera_name = receivedJson.camera_data.name
                sendJson.data.result.total_result = total_result
                sendJson.data.result.standard_result.append({
                    'standard_id':getattr(receivedJson.camera_data.standards[0],'standard_id',''),
                    'label':receivedJson.camera_data.standards[0].label,
                    'type':receivedJson.camera_data.standards[0].type,
                    'standard':receivedJson.camera_data.standards[0].standard,
                    'threshold':receivedJson.camera_data.standards[0].threshold,
                    'scores':float(max_scores),
                    'algo_result':algo_result,
                    'detect_result':total_result,
                    
                })
                
                send = json.dumps(sendJson, default=lambda x: x.__dict__).encode('utf8')
                print(f'================send json================')
                pprint.pprint(sendJson)
                return send
            
    
def model_init(yolo_model_handian_weights,yolo_model_gap_weights):
    global yolo_model_handian, yolo_model_gap
    try:
        yolo_model_handian = get_model(yolo_model_handian_weights)
        yolo_model_gap = get_model_obb(yolo_model_gap_weights)
    except Exception as e:
        print('模型加载失败，检查模型路径')
        print(e)

def flask_init():
    host = "127.0.0.1"
    port = 8080
    # app.debug=True
    print(app.url_map)
    app.run(host=host, port=port, threaded=False)
    
if __name__ == '__main__':
    model_init(yolo_model_handian_weights,
               yolo_model_gap_weights)
    flask_init()